package com.alibaba.otter.node.etl.load.loader.mq.rabbit.consume;

import com.alibaba.otter.node.etl.load.loader.mq.rabbit.MqConstants;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.MqUtils;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.callback.ConnectedCallback;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.callback.ConsumeFailureCallback;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.callback.DisconnectedCallback;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.exception.MqConsumeException;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.exception.MqSerializeException;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.model.Message;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.model.MessageContext;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.strategy.FailureAction;
import com.alibaba.otter.node.etl.load.loader.mq.rabbit.strategy.MqExceptionDispatcher;
import com.google.common.collect.Maps;
import com.rabbitmq.client.AMQP;
import com.rabbitmq.client.Envelope;
import com.rabbitmq.client.ShutdownSignalException;
import com.rabbitmq.client.impl.recovery.AutorecoveringChannel;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.IntStream;

@Slf4j
public class BlockingQueueConsumer extends AbstractConsumer {

    public BlockingQueueConsumer(AutorecoveringChannel channel, String namespace) {
        this.channel = channel;
        this.namespace = namespace;
    }

    @Setter
    private int threadCount = MqConstants.DEFAULT_CONSUME_THREAD_COUNT;

    private String namespace;

    private BlockingQueue<MessageContext> blockingQueue = new LinkedBlockingQueue<>(MqConstants.MAX_PREFETCH_SIZE);

    private volatile AtomicBoolean started = new AtomicBoolean(false);

    private ExecutorService executorService;

    @Setter
    private DefaultHandlerInvoker handlerInvoker;

    @Setter
    private AutorecoveringChannel channel;

    @Setter
    private ConnectedCallback connectedCallback;

    @Setter
    private DisconnectedCallback disconnectedCallback;

    @Setter
    private List<ConsumeFailureCallback> failureCallbacks;

    public void setFailureAction(FailureAction failureAction) {
        this.failureAction = failureAction;
        dispatcher.setAction(failureAction);
    }

    private FailureAction failureAction;

    private String queue;

    private Runnable task = () -> {
        while (started.get()) {
            // take message from queue, if fail, retry.
            MessageContext context = getMessageFromQueue();
            // invoke handler
            invokeHandler(context);
        }
    };

    private MessageContext getMessageFromQueue() {
        MessageContext context = null;
        while (context == null) {
            try {
                context = blockingQueue.take();
            } catch (InterruptedException e) {
                String errMsg = "failed to get message from local blocking queue, retry.";
                log.error(errMsg, e);
            }
        }
        return context;
    }

    private void nack(MessageContext ctx) {
        if (ctx == null || ctx.getEnvelope() == null) {
            return;
        }
        long deliveryTag = ctx.getEnvelope().getDeliveryTag();
        try {
            channel.basicNack(deliveryTag, false, true);
        } catch (IOException e) {
            log.error("failed to nack message: " + deliveryTag, e);
        }
    }

    public void listen(String queue) {
        this.queue = queue;
        this.disconnectedCallback.setChannel(channel);
        started.set(true);
        executorService = Executors.newFixedThreadPool(
                threadCount,
                new CustomNameThreadFactory(MqConstants.CONSUMER_THREAD_NAME)
        );
        IntStream.range(0, threadCount).forEach(i -> executorService.execute(task));
    }

    private void invokeHandler(MessageContext context) {
        try {
            handlerInvoker.invoke(context);
        } catch (MqSerializeException e) {
            log.error("反序列化失败. MessageContext: " + context, e);
            if (failureAction == null) {
                failureCallbacks.forEach(c -> c.onDeserializeFailure(e.getMessage(), e.getCause()));
                MqUtils.ack(context);
            } else {
                dispatcher.dispatch(new MqConsumeException(e.getMessage(), context, e));
            }
        } catch (MqConsumeException e) {
            log.error("消息消费失败. MessageContext: " + context, e);
            if (failureAction == null) {
                failureCallbacks.forEach(c -> c.onBizFailure(e.getMessage(), e.getCause()));
                MqUtils.ack(context);
            } else {
                dispatcher.dispatch(e);
            }
        } catch (Exception e) {
            log.error("调用handler失败. MessageContext: " + context, e);
            if (failureAction == null) {
                failureCallbacks.forEach(c -> c.onBizFailure(e.getMessage(), e.getCause()));
                MqUtils.ack(context);
            } else {
                dispatcher.dispatch(new MqConsumeException(e.getMessage(), context, e));
            }
        }
    }

    private MqExceptionDispatcher dispatcher = new MqExceptionDispatcher();

    @Override
    public void handleDelivery(String consumerTag,
                               Envelope envelope,
                               AMQP.BasicProperties properties,
                               byte[] body) {
        Message message = new Message(body, properties.getHeaders(), properties);
        MessageContext context = new MessageContext(
                message, envelope, null, channel, handlerInvoker, namespace, queue, Maps.newConcurrentMap()
        );
        try {
            blockingQueue.put(context);
        } catch (InterruptedException e) {
            log.warn("线程中断.", e);
            nack(context);
            Thread.currentThread().interrupt();
        }
    }

    @Override
    public void handleConsumeOk(String consumerTag) {
        connectedCallback.connectedCallback(channel);
        started.set(true);
    }

    @Override
    public void handleShutdownSignal(String consumerTag, ShutdownSignalException sig) {
        disconnectedCallback.onDisconnected(sig);
        blockingQueue.clear();
        started.set(false);
        listen(queue);
    }

    @Override
    public void handleCancel(String consumerTag) {
        log.info("handle cancel, restarting: " + consumerTag);
        channel.notifyListeners();
    }

    @Override
    public void handleRecoverOk(String consumerTag) {
        log.info("handleRecoverOk, consumerTag: " + consumerTag);
    }

    public void stop() {
        executorService.shutdown();
        blockingQueue.clear();
        started.set(false);
        try {
            if (channel.isOpen()) {
                channel.basicCancel(super.consumerTag);
                channel.close();
            }
        } catch (IOException | TimeoutException e) {
            log.error(e.getMessage(), e);
        }
    }
}
